{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_all(121150643,\n",
    "#          rpt_only=True,\n",
    "#          large=True,\n",
    "#          bkspace=5.0,\n",
    "#          abs_xlim='3p',\n",
    "#          abs_offset=6,\n",
    "#          abs_xlim=(1300, 2500),\n",
    "#          abs_ylim='3%',\n",
    "#          abs_ylim=(0.975, 1.01),\n",
    "#          item_no=5\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run at startup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "from astronet.preprocess import generate_input_records\n",
    "from astronet.preprocess import preprocess\n",
    "from light_curve_util import keplersplinev2\n",
    "from light_curve_util import tess_io\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tce_tables = {}\n",
    "\n",
    "tce_files = [\n",
    "    '/mnt/tess/astronet/tces-vetting-all.csv',\n",
    "    '/mnt/tess/astronet/tces-vetting-v5-toi-train.csv',\n",
    "    '/mnt/tess/astronet/tces-vetting-v5-toi-val.csv',\n",
    "    '/mnt/tess/astronet/tces-vetting-v5-toi-test.csv',\n",
    "]\n",
    "\n",
    "\n",
    "def get_tce(tic, item_no):\n",
    "    for f in tce_files:\n",
    "        if f not in tce_tables:\n",
    "            tce_tables[f] = pd.read_csv(f, header=0)\n",
    "        tce_table = tce_tables[f]\n",
    "        tce = tce_table[tce_table.tic_id == tic]\n",
    "        if not len(tce):\n",
    "            continue\n",
    "        if len(tce) > 1:\n",
    "            tce = tce[tce.index == tce.index.values[item_no]]\n",
    "            if 'Source' in tce:\n",
    "                print('Source:', 'BLS' if tce.Source.values.item() == 2 else 'TEV')\n",
    "        return tce, f\n",
    "    raise ValueError(f'no TCE data for {tic}')\n",
    "\n",
    "\n",
    "def plot_all(tic,\n",
    "             large=False,\n",
    "             dpi=100,\n",
    "             bkspace=None,\n",
    "             abs_xlim=None,\n",
    "             abs_offset=0,\n",
    "             abs_ylim=None,\n",
    "             item_no=0):\n",
    "    tess_data_dir = '/mnt/tess/lc-v'\n",
    "    reports_dir = f'/mnt/tess/rpt/png-v'\n",
    "\n",
    "    tce, tces_file = get_tce(tic, item_no)\n",
    "    generate_input_records.FLAGS = generate_input_records.parser.parse_args([\n",
    "      '--tess_data_dir', tess_data_dir,\n",
    "      '--output_dir', '/dev/null',\n",
    "      '--input_tce_csv_file', tces_file,\n",
    "      '--vetting_features', 'y',\n",
    "    ])\n",
    "\n",
    "\n",
    "    b = plt.get_cmap('tab20')(0)\n",
    "    b2 = plt.get_cmap('tab20')(1)\n",
    "    o = plt.get_cmap('tab20')(2)\n",
    "    o2 = plt.get_cmap('tab20')(3)\n",
    "    g = plt.get_cmap('tab20')(4)\n",
    "    r = plt.get_cmap('tab20')(6)\n",
    "    n = plt.get_cmap('tab20')(8)\n",
    "    k = plt.get_cmap('tab20')(14)\n",
    "\n",
    "    plotrows = 7\n",
    "    plotcols = 3\n",
    "    \n",
    "    if large:\n",
    "        fsize = (16, 9)\n",
    "    else:\n",
    "        fsize = (16, 4 * (plotrows + 1))\n",
    "        plt.figure(figsize=fsize, dpi=dpi)\n",
    "    \n",
    "    period = tce.Period.values.item()\n",
    "    epoc = tce.Epoc.values.item()\n",
    "    duration = tce.Duration.values.item()\n",
    "\n",
    "    print(f'Epoc: {epoc}\\nPeriod: {period}\\nDuration: {duration}')\n",
    "    \n",
    "    def config_abs_plot(title):\n",
    "        plt.legend()\n",
    "        plt.title(title)\n",
    "        if abs_xlim:\n",
    "            if abs_xlim == '3p':\n",
    "                minx = min(td) + abs_offset * period\n",
    "                maxx = minx + 3.5 * period\n",
    "                plt.xlim(minx, maxx)\n",
    "            else:\n",
    "                plt.xlim(*abs_xlim)\n",
    "        if abs_ylim:\n",
    "            if abs_ylim == '3%':\n",
    "                miny = np.percentile(fs[~np.isnan(fs)], 3)\n",
    "                maxy = np.percentile(fs[~np.isnan(fs)], 97)\n",
    "                plt.ylim(miny, maxy)\n",
    "            else:\n",
    "                plt.ylim(*abs_ylim)\n",
    "                \n",
    "    nplotted = 0\n",
    "    def splt(c=1):\n",
    "        nonlocal nplotted\n",
    "        if large:\n",
    "            plt.figure(figsize=fsize, dpi=dpi)\n",
    "        else:\n",
    "            plt.subplot(plotrows, plotcols // c, (nplotted // c) + 1)\n",
    "        nplotted += c\n",
    "    \n",
    "    \n",
    "    file_names = tess_io.tess_filenames(tic, tess_data_dir)\n",
    "    f = fits.open(file_names)\n",
    "    \n",
    "    td = f[1].data[\"TIME\"]\n",
    "    if \"KSPSAP_FLUX\" in f[1].data:\n",
    "      fd = f[1].data[\"KSPSAP_FLUX\"]\n",
    "    else:\n",
    "      fd = []\n",
    "    fs = f[1].data[\"SAP_FLUX\"]\n",
    "    \n",
    "    splt()\n",
    "    plt.plot(td, fs, '-', alpha=0.6, color=g, label='SAP')\n",
    "    config_abs_plot('fits data')\n",
    "    \n",
    "    \n",
    "    splt()\n",
    "    ut, uf, ap = preprocess.read_and_process_light_curve(\n",
    "        tic, tess_data_dir, 'SAP_FLUX',\n",
    "        {\n",
    "            's': 'SAP_FLUX_SML',\n",
    "            'm': 'SAP_FLUX_MID',\n",
    "            'l': 'SAP_FLUX_LAG',\n",
    "        },\n",
    "    )\n",
    "    input_mask = preprocess.get_spline_mask(ut, period, epoc, duration)\n",
    "    sf, mdata = keplersplinev2.choosekeplersplinev2(\n",
    "        ut, uf, input_mask=input_mask, return_metadata=True,\n",
    "        fixed_bkspace=bkspace,\n",
    "    )\n",
    "\n",
    "    plt.plot(ut, uf, '-', alpha=0.6, color=g, label='SAP')\n",
    "    plt.plot(ut[input_mask], sf[input_mask], 'x', markersize=3, color=k, label='spline (OOT)')\n",
    "    plt.plot(ut[~input_mask], sf[~input_mask], 'o', markersize=3, color=o, label='spline')\n",
    "    config_abs_plot(f'raw | bkspace: {mdata.bkspace}')\n",
    "    \n",
    "    \n",
    "    ut, nf, fm = preprocess.detrend_and_filter(tic, ut, uf, period, epoc, duration, bkspace)\n",
    "    sft, sff, sfn, sftm = preprocess.phase_fold_and_sort_light_curve(ut, nf, input_mask, period, epoc)\n",
    "\n",
    "    splt()\n",
    "    plt.plot(sft, sff, 'o', markersize=3, alpha=0.6, c=o, label='spline')\n",
    "    sff_filtered = np.where((sff > 1.5) | (sff < -0.5), 0, sff)\n",
    "    if len(sff_filtered):\n",
    "        mask = np.where(sftm, 1, min(sff_filtered))\n",
    "        plt.plot(sft, mask, '-', markersize=1, alpha=0.6, c=r, label='OOT')\n",
    "        title = f'phase folded ({int(max(sfn) + 1)} folds) / filtered'\n",
    "    else:\n",
    "        if len(sfn):\n",
    "            title = f'phase folded ({int(max(sfn) + 1)} folds) | WARNING: filtering removed all data'\n",
    "        else:\n",
    "            title = f'phase folded | WARNING: filtering removed all data'\n",
    "    plt.legend()\n",
    "    plt.title(title)\n",
    "\n",
    "\n",
    "    splt()\n",
    "    if ap:\n",
    "        uf_s = ap['s'][1]\n",
    "        uf_m = ap['m'][1]\n",
    "        uf_l = ap['l'][1]\n",
    "\n",
    "        _, nf_s, _ = preprocess.detrend_and_filter(tic, ut, uf_s, period, epoc, duration, bkspace)\n",
    "        sft_s, sff_s, _, _ = preprocess.phase_fold_and_sort_light_curve(ut, nf_s, input_mask, period, epoc)\n",
    "        _, nf_m, _ = preprocess.detrend_and_filter(tic, ut, uf_m, period, epoc, duration, bkspace)\n",
    "        sft_m, sff_m, _, _ = preprocess.phase_fold_and_sort_light_curve(ut, nf_m, input_mask, period, epoc)\n",
    "        _, nf_l, _ = preprocess.detrend_and_filter(tic, ut, uf_m, period, epoc, duration, bkspace)\n",
    "        sft_l, sff_l, _, _ = preprocess.phase_fold_and_sort_light_curve(ut, nf_m, input_mask, period, epoc)\n",
    "\n",
    "        plt.plot(sft_s, sff_s, '.', markersize=1, alpha=0.6, label='sm')\n",
    "        plt.plot(sft_m, sff_m, '.', markersize=1, alpha=0.6, label='med')\n",
    "        plt.plot(sft_l, sff_l, '.', markersize=1, alpha=0.6, label='lg')\n",
    "        plt.legend()\n",
    "        plt.title('multi-aperture')\n",
    "\n",
    "\n",
    "    row = list(tce.iterrows())[0][1]\n",
    "\n",
    "    ex = generate_input_records._process_tce(row, bkspace)\n",
    "    \n",
    "    if ex.features.feature['star_rad_present'].int64_list.value[0] > 0:\n",
    "        print(f'star_rad: {ex.features.feature[\"star_rad\"].float_list.value[0]}')\n",
    "    else:\n",
    "        print('no star_rad')\n",
    "    print(f'Transit_Depth: {ex.features.feature[\"Transit_Depth\"].float_list.value[0]}')\n",
    "    print(f'local_scale: {ex.features.feature[\"local_scale\"].float_list.value[0]}')\n",
    "    \n",
    "    \n",
    "    def plot_w_scatter(view, std, title, mask=None, mask_is_filter=False):\n",
    "        splt()\n",
    "        view = np.array(ex.features.feature[view].float_list.value)\n",
    "        std = np.array(ex.features.feature[std].float_list.value)\n",
    "        if mask:\n",
    "            msk = -np.array(ex.features.feature[mask].float_list.value)\n",
    "            if mask_is_filter:\n",
    "                view = view[msk < 0]\n",
    "                std = std[msk < 0]\n",
    "        plt.plot(std, color=o)\n",
    "        plt.plot(view, color=n)\n",
    "        if mask:\n",
    "            plt.plot(msk, color=r, alpha=0.5, linestyle='--')\n",
    "        plt.title(title)\n",
    "        \n",
    "    def plot_norm(view, title):\n",
    "        splt()\n",
    "        plt.plot(ex.features.feature[view].float_list.value, color=n)\n",
    "        plt.ylim(-1, 1)\n",
    "        plt.title(title)    \n",
    "        \n",
    "        \n",
    "    def plot_segments(view, tag, num_chans=14, min_chans=None, max_chans=None):\n",
    "        splt(c=3)\n",
    "        img = ex.features.feature[view].float_list.value\n",
    "        img = np.reshape(img, (-1, num_chans))[:, min_chans:max_chans]\n",
    "        n_transits = img.shape[1] // 2\n",
    "        for i in range(n_transits):\n",
    "            view = img[:, 2 * i]\n",
    "            mask = img[:, 2 * i + 1] > 0\n",
    "            plt.plot(np.where(mask, view, np.nan), marker='.')\n",
    "        plt.title(f'{n_transits} sample segments, densest first, ties broken at random {tag}')\n",
    "        \n",
    "    splt()\n",
    "    plt.plot(ex.features.feature['local_aperture_s'].float_list.value, label='small')\n",
    "    plt.plot(ex.features.feature['local_aperture_m'].float_list.value, label='med')\n",
    "    plt.plot(ex.features.feature['local_aperture_l'].float_list.value, label='large')\n",
    "    plt.title('multi-aperture')\n",
    "    plt.legend()\n",
    "    \n",
    "    splt()\n",
    "\n",
    "    plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3')\n",
    "    plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0')\n",
    "\n",
    "    sec_phase = ex.features.feature['secondary_phase'].float_list.value[0]\n",
    "    loc_scale = ex.features.feature['local_scale'].float_list.value[0]\n",
    "    sec_scale = ex.features.feature['secondary_scale'].float_list.value[0]\n",
    "    plot_w_scatter('global_view', 'global_std', 'global', 'global_transit_mask')\n",
    "    plot_w_scatter('local_view', 'local_std', f'local / {loc_scale:0.3}')\n",
    "    plot_w_scatter('secondary_view', 'secondary_std', f'secondary ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "    plot_w_scatter('local_view_half_period', 'local_view_half_period_std', 'local half period')\n",
    "    plot_w_scatter('local_view_even', 'local_std_even', 'local even', 'local_mask_even', True)\n",
    "    plot_w_scatter('local_view_odd', 'local_std_odd', 'local odd', 'local_mask_odd', True)\n",
    "    \n",
    "    splt()\n",
    "\n",
    "    plot_segments('sample_segments_local_view', 'odd', num_chans=16, min_chans=0, max_chans=8)\n",
    "    plot_segments('sample_segments_local_view', 'even', num_chans=16, min_chans=8, max_chans=16)\n",
    "\n",
    "    try:\n",
    "        im = plt.imread(f'{reports_dir}/{tic}.page1.png')\n",
    "        plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "        plt.axis('off')\n",
    "        _ = plt.imshow(im)\n",
    "        try:\n",
    "            im = plt.imread(f'{reports_dir}/{tic}.page3.png')\n",
    "            plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "            plt.axis('off')\n",
    "            _ = plt.imshow(im)\n",
    "        except FileNotFoundError:\n",
    "            print('-- no page 3 --')\n",
    "    except FileNotFoundError:\n",
    "        print('-- no report file --')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
